package com.dman.multi.proxy;

import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.net.InetSocketAddress;
import java.net.SocketException;
import java.net.URL;
import java.nio.ByteBuffer;
import java.nio.channels.SocketChannel;
import java.util.HashMap;
import java.util.Map;

import org.apache.pivot.util.Base64;

public class Socks5ProxyHandler implements ProxyHandler {
	
	protected static final String SOCKS_CHARSET = "US-ASCII";
	protected Map<Integer, String> ecodeMessage = null;
	
	protected static final int SOCKS_VER = 5;
	
	protected static final int NO_AUTH = 0;
	protected static final int GSSAPI = 1;
	protected static final int USER_PASS = 2;
	protected static final int NO_SUPPORTED_METH = 0xFF;
	
	protected static final int TCP_CONNECT = 1;
	protected static final int TCP_BIND = 2;
	protected static final int UDP = 3;
	
	protected static final int IPV4 = 1;
	protected static final int DOMAIN_NAME = 3;
	protected static final int IPV6 = 4;
	
	public Socks5ProxyHandler() {
		ecodeMessage = new HashMap<Integer, String>();
	    ecodeMessage.put(0x00, "SOCKS: request granted");
	    ecodeMessage.put(0x01, "SOCKS: general failure");
	    ecodeMessage.put(0x02, "SOCKS: connection not allowed by ruleset");
	    ecodeMessage.put(0x03, "SOCKS: network unreachable");
	    ecodeMessage.put(0x04, "SOCKS: host unreachable");
	    ecodeMessage.put(0x05, "SOCKS: connection refused by destination host");
	    ecodeMessage.put(0x06, "SOCKS: TTL expired");
	    ecodeMessage.put(0x07, "SOCKS: command not supported / protocol error");
	    ecodeMessage.put(0x08, "SOCKS: address type not supported");
	}
	
	@Override
	public boolean connect(SocketChannel sockChannel, ProxyAdaptor proxyAuth,
			URL url) throws IOException {
		int version = getSocksVersion(sockChannel, proxyAuth);
		byte authMethod = (byte) version;
		byte versionNumber = (byte) (version >> 8);
		boolean isValid = authenticate(sockChannel, proxyAuth, authMethod);
		if(isValid){
			// if authenticated check if address can be resolved
			byte [] address = resolveAddress(url);
			
			// now create request and send it
			byte [] reqBytes = new byte [address.length + 6];
			int index = 0;
			reqBytes[index++] = versionNumber;
			reqBytes[index++] = TCP_CONNECT;
			reqBytes[index++] = 0;
			System.arraycopy(address, 0, reqBytes, index, address.length);
			index += address.length;
			sockChannel.write(ByteBuffer.wrap(reqBytes));
			
			// now read the response
			byte [] resBytes = new byte [4];
			sockChannel.read(ByteBuffer.wrap(resBytes));
			if(resBytes[1] != 0){
				throw new IOException(ecodeMessage.get(resBytes[1]));
			}
			
			return true;
		}
		
		return false;
	}
	
	public byte [] resolveAddress(URL url) throws UnsupportedEncodingException{
		
		// get the host and port and create a socket address
		String host = url.getHost();
		int port = url.getPort();
		if(port < -1) 
			port = url.getDefaultPort();			
		InetSocketAddress isa = new InetSocketAddress(host, port);
		
		// assume an unresolved address
		byte [] hostBytes = host.getBytes(SOCKS_CHARSET);
		byte [] address = new byte [hostBytes.length + 4];
		int index = 0;
		address[index++] = DOMAIN_NAME;
		address[index++] = (byte) hostBytes.length;
		System.arraycopy(hostBytes, 0, address, index, hostBytes.length);
		index += hostBytes.length;
		
		// if you had resolved the address modify the address
		if(!isa.isUnresolved()){
			index -= hostBytes.length;
			byte[] inetAddr = isa.getAddress().getAddress();
			address = new byte [inetAddr.length + 3];
			System.arraycopy(inetAddr, 0, address, index, inetAddr.length);
			address[0] = IPV4;
			if(inetAddr.length == 16)
				address[0] = IPV6;
			index += inetAddr.length;
		}
		
		// write the port
		address[index++] = (byte) (port / 0x100);
		address[index++] = (byte) (port % 0x100);
		
		return address;
	}
	
	public int getSocksVersion(SocketChannel sockChannel, ProxyAdaptor proxyAuth) throws IOException{
		byte [] reqBytes = null;
		byte [] resBytes = new byte [2];
		int version = 0;
		
		/*
		 * if the proxy authentication is null only NO_AUTH is supported, otherwise
		 * NO_AUTH as well as USER_PASS is supported.
		 */
		if(proxyAuth.getAuth() != null){
			reqBytes = new byte [3];
			reqBytes[0] = SOCKS_VER;
			reqBytes[1] = 1;
			reqBytes[2] = NO_AUTH;
		} else {
			reqBytes = new byte [4];
			reqBytes[0] = SOCKS_VER;
			reqBytes[1] = 2;
			reqBytes[2] = NO_AUTH;
			reqBytes[3] = USER_PASS;
		}
		sockChannel.connect(proxyAuth.getProxy().address());
		sockChannel.write(ByteBuffer.wrap(reqBytes));
		sockChannel.read(ByteBuffer.wrap(resBytes));
		
		// return the SOCKS version and supported authentication method
		version = resBytes[0];
		version = version << 8;
		version += resBytes[1];
		
		return version;
	}
	
	public boolean authenticate(SocketChannel sockChannel, ProxyAdaptor proxyAuth, byte authMethod) throws IOException{
		
		if(authMethod == NO_SUPPORTED_METH)
			throw new SocketException("Unsupported authentication mechanism");
		
		if(authMethod == NO_AUTH)
			return true;
		
		if(authMethod == USER_PASS){
			// get the authentication details
			String auth = proxyAuth.getAuth();
			String actualAuth = new String(Base64.decode(auth), SOCKS_CHARSET);
			String array [] = actualAuth.split(":");
			
			// create the authentication byte array
			byte [] byteArray = new byte[actualAuth.length() + 2];
			int index = 0;
			byteArray[index++] = 1;
			byteArray[index++] = (byte) array[0].length();
			byte [] strBytes = array[0].getBytes(SOCKS_CHARSET);
			System.arraycopy(strBytes, 0, byteArray, index, array[0].length());
			index += array[0].length();
			byteArray[index++] = (byte) array[1].length();
			strBytes = array[1].getBytes();
			System.arraycopy(strBytes, 0, byteArray, index, array[1].length());
			index += array[1].length();
			
			// write the request and read response
			byte [] resArray = new byte [2];
			sockChannel.write(ByteBuffer.wrap(byteArray));
			sockChannel.read(ByteBuffer.wrap(resArray));
			
			// verify if authenticated
			if(resArray[1] == 0)
				return true;
		}
		
		return false;
	}
	
	@Override
	public String identifyProtocol(URL url, ProxyAdaptor proxyAuth)
			throws IOException {
		return url.getProtocol();
	}
}
